import cv2
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F


class ImageProcessing(nn.Module):

    # ---------------------- 图像预处理 ----------------------
    @staticmethod
    def preprocess_image(img, is_ir=False, size=(320, 180)):
        gray = cv2.cvtColor(img, cv2.COLOR_BGR2GRAY) if len(img.shape) == 3 else img
        if is_ir:
            gray = cv2.createCLAHE(clipLimit=2.0, tileGridSize=(8, 8)).apply(gray)
        h, w = gray.shape
        ratio = min(size[0] / w, size[1] / h)
        new_w, new_h = int(w * ratio), int(h * ratio)
        resized = cv2.resize(gray, (new_w, new_h), interpolation=cv2.INTER_AREA)
        pad_w = size[0] - new_w
        pad_h = size[1] - new_h
        padded = cv2.copyMakeBorder(resized, pad_h // 2, pad_h - pad_h // 2, pad_w // 2, pad_w - pad_w // 2,
                                    cv2.BORDER_CONSTANT)
        return padded.astype(np.float32) / 255.0, (h, w, new_w, new_h, pad_w // 2, pad_h // 2)

    # ---------------------- 特征提取 ----------------------
    @staticmethod
    def extract_keypoints(img_tensor, model, device, max_keypoints=1000):
        with torch.no_grad():
            output = model(img_tensor.unsqueeze(0).unsqueeze(0).to(device))
        scores = output['scores'][0].cpu().numpy()
        descriptors = output['descriptors'][0].cpu().numpy().T
        h, w = img_tensor.shape
        y, x = np.where(scores > 0.01)
        scores = scores[y, x]
        keypoints = np.stack([x, y], axis=1)
        indices = np.argsort(scores)[::-1][:max_keypoints]
        return keypoints[indices], descriptors[indices]
